#include <fstream>
#include <string>
#include <filesystem>
#include "itkImageSeriesReader.h"
#include "itkGDCMImageIO.h"
#include "itkGDCMSeriesFileNames.h"

namespace lzw
{
    using PixelType = unsigned short;
    constexpr unsigned int Dimension = 3;
    using ImageType = itk::Image<PixelType, Dimension>;
    using ReaderType = itk::ImageSeriesReader<ImageType>;
    using ImageIOType = itk::GDCMImageIO;
    using NamesGeneratorType = itk::GDCMSeriesFileNames;

    // Function to create an MHD file
    void CreateMHDFile(const std::string& mhdFileName, const std::string& rawFileName, const ImageType::Pointer& image)
    {
        std::filesystem::path mhdPath = std::filesystem::path(mhdFileName).parent_path();
        if (!mhdPath.empty() && !std::filesystem::exists(mhdPath))
        {
            std::filesystem::create_directories(mhdPath);
        }

        std::ofstream mhdFile(mhdFileName);
        if (!mhdFile.is_open())
        {
            std::cerr << "Error opening MHD file: " << mhdFileName << std::endl;
            return;
        }

        ImageType::SizeType size = image->GetLargestPossibleRegion().GetSize();
        ImageType::SpacingType spacing = image->GetSpacing();

        mhdFile << "ObjectType = Image\n";
        mhdFile << "NDims = 3\n";
        mhdFile << "DimSize = " << size[0] << " " << size[1] << " " << size[2] << "\n";
        mhdFile << "ElementType = MET_USHORT\n";
        mhdFile << "ElementSpacing = " << spacing[0] << " " << spacing[1] << " " << spacing[2] << "\n";
        mhdFile << "ElementByteOrderMSB = False\n";
        mhdFile << "ElementDataFile = " << rawFileName << "\n";

        mhdFile.close();

        std::cout << "Conversion successful: " << mhdFileName << std::endl;
    }

    ImageType::Pointer CreateRawFile(const std::string& dicomDirectory, const std::string& rawFileName)
    {
        std::filesystem::path rawPath = std::filesystem::path(rawFileName).parent_path();
        if (!rawPath.empty() && !std::filesystem::exists(rawPath))
        {
            std::filesystem::create_directories(rawPath);
        }

        NamesGeneratorType::Pointer nameGenerator = NamesGeneratorType::New();
        nameGenerator->SetDirectory(dicomDirectory);

        const auto seriesUIDs = nameGenerator->GetSeriesUIDs();
        if (seriesUIDs.empty())
        {
            std::cerr << "No DICOM series found in directory: " << dicomDirectory << std::endl;
            return nullptr;
        }

        std::string seriesUID = seriesUIDs.front();
        std::cout << "Reading series: " << seriesUID << std::endl;

        const auto fileNames = nameGenerator->GetFileNames(seriesUID);

        ReaderType::Pointer reader = ReaderType::New();
        ImageIOType::Pointer dicomIO = ImageIOType::New();

        reader->SetImageIO(dicomIO);
        reader->SetFileNames(fileNames);

        try
        {
            reader->Update();
        }
        catch (itk::ExceptionObject& ex)
        {
            std::cerr << "Exception caught: " << ex << std::endl;
            return nullptr;
        }

        // Get pointer to image data
        ImageType::Pointer image = reader->GetOutput();
        const PixelType* buffer = image->GetBufferPointer();
        size_t totalSize = image->GetBufferedRegion().GetNumberOfPixels();

        // Write buffer to RAW file
        std::ofstream outFile(rawFileName, std::ios::binary);
        if (!outFile)
        {
            std::cerr << "Error opening output file: " << rawFileName << std::endl;
            return nullptr;
        }
        outFile.write(reinterpret_cast<const char*>(buffer), totalSize * sizeof(PixelType));
        outFile.close();

        std::cout << "Conversion successful: " << rawFileName << std::endl;
        return image;
    }

}